
# Copyright (c) 2022 Presto Labs Pte. Ltd.
# Author: zijin

import matplotlib

matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.dates as mdates

import os
import datetime
import tempfile
import traceback
import warnings

import numpy as np
import pandas as pd
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
from coin.base.mail_util import send_mail_with_attachments
from python.coin.experimental.yuxuan.model.regression3 import get_trading_dates
from python.coin.experimental.zijin.retrain_model_weekly import split_list
from coin2.service.account.account_info_service_pb2 import QueryBalanceIntervalHistoriesRequestProto
from cc.appcoin2.strategy.basis_strat.support.daily_search.slack_helper import SlackHelper
from google.protobuf.json_format import MessageToDict
from xunkemgmt_client.client.util.query_util import query_accounts
from coin.base.query_util import (query_exchange_rates, query_pta)
from xunkemgmt_client.client.api_client import XunkemgmtClient
from coin.base.datetime_util import to_timestamp_int
from coin.proto.coin_request_pb2 import (ApplicationRequestProto,
                                         StrategyRequestProto)


from absl import app, flags
FLAGS = flags.FLAGS


def daterange(start_date, end_date):
  for n in range(int((end_date - start_date).days)):
    yield start_date + datetime.timedelta(n)


def get_single_rates(quotes, single_date):
  return {quote: rate for quote, rate in zip(
      quotes, query_exchange_rates(quotes, "USD", single_date.date()))}


def get_rate_per_symbol_day(balance_df, start_date, end_date):
  rate_dict = dict()
  start_date = datetime.datetime.strptime(start_date, '%Y%m%d')
  end_date = datetime.datetime.strptime(end_date, '%Y%m%d')
  quotes = [list(bdict.keys()) for bdict in balance_df['balance']]
  quotes = set([item for sublist in quotes for item in sublist])
  with ProcessPoolExecutor(max_workers=FLAGS.max_workers) as executor:
    futures = {}
    for single_date in daterange(start_date, end_date):
      fut = executor.submit(get_single_rates, quotes, single_date)
      futures[fut] = single_date
    for f in as_completed(futures.keys()):
      try:
        rate_dict[futures[f].date()] = f.result()
      except Exception:
        import traceback
        traceback.print_exc()
  return rate_dict


def get_historical_account_history(start_date, end_date):
  business_unit = 'Coin'
  account_df = query_accounts(cache_clear=FLAGS.cache_clear)
  start_date = datetime.datetime.strptime(start_date, '%Y%m%d')
  end_date = datetime.datetime.strptime(end_date, '%Y%m%d')

  with XunkemgmtClient() as api_client:
    balance_request_dict = {
        'start_timestamp': to_timestamp_int(start_date),
        'end_timestamp': to_timestamp_int(end_date),
        'business_units': [business_unit]
    }
    balance_request = QueryBalanceIntervalHistoriesRequestProto(**balance_request_dict)
    responses = api_client.query_balance_interval_histories(balance_request)
    responses = responses.histories
    balance_list = []
    for response in responses:
      response = {
          **MessageToDict(response.balance_info, preserving_proto_field_name=True),
          "each_balance": [
              MessageToDict(balance, preserving_proto_field_name=True)
              for balance in response.each_balance]
      }
      balance_list.append(response)
    balance_df = pd.DataFrame(balance_list)

    app_request = ApplicationRequestProto()
    responses = api_client.query_applications(app_request)
    app_list = []
    for response in responses:
      response = MessageToDict(response, preserving_proto_field_name=True)
      app_list.append(response)
    application_df = pd.DataFrame(app_list)

    strat_request = StrategyRequestProto()
    responses = api_client.query_strategies(strat_request)
    strategy_list = []
    for response in responses:
      response = MessageToDict(response, preserving_proto_field_name=True)
      strategy_list.append(response)
    strategy_df = pd.DataFrame(strategy_list)

  strategy_df = pd.merge(application_df, strategy_df, on='strategy_id', suffixes=('', '_y'))
  account_df = account_df[['account_id', 'market_type', 'exchange', 'api_version']]
  balance_df['balance'] = balance_df['each_balance'].apply(
        lambda x: {b['currency']: b['total'] for b in x})
  balance_df['trading_date'] = balance_df['trading_date'].apply(
        lambda x: datetime.datetime.strptime(str(x), '%Y%m%d').date())
  balance_df = balance_df[['account_id', 'trading_date', 'balance']]
  strategy_df = strategy_df[['strategy_name', 'strategy_group', 'account_id']]
  strategy_df = pd.merge(account_df, strategy_df, on='account_id')
  merged_balance_df = pd.merge(strategy_df, balance_df, on='account_id')
  return merged_balance_df


def get_historical_pta(start_date, end_date):
  dates = get_trading_dates('-'.join([start_date, end_date]))
  date_groups = split_list(dates, FLAGS.chunk_size)
  dfs = []
  with ThreadPoolExecutor(max_workers=FLAGS.max_workers) as executor:
    futures = {}
    for dates in date_groups:
      start_date = datetime.datetime.strptime(dates[0], '%Y%m%d')
      end_date = datetime.datetime.strptime(dates[-1], '%Y%m%d')
      fut = executor.submit(query_pta,
                            start_date,
                            end_date,
                            business_units=['Coin'],
                            pta_type='STRATEGY')
      futures[fut] = start_date
    for f in as_completed(futures.keys()):
      try:
        pta_data = f.result()
        dfs.append(pta_data.rename(
            columns={
                'fill_pq_in_usd': 'volume',
                'pnl_mark_in_usd': 'pnl'}))
        print(f'{futures[f]} done')
      except Exception:
        import traceback
        traceback.print_exc()
  pta_details = pd.concat(dfs)[['strategy_name', 'strategy_group',
                                'trading_date', 'pnl', 'volume']].fillna(0)
  pta_details = pta_details.loc[(pta_details[['pnl', 'volume']] != 0).any(axis=1), :]
  return pta_details


def calculate_balance_in_usd(bdict, rate_list):
  rate_array = np.array([rate_list[quote] for quote in bdict.keys()], dtype=np.float64)
  val_array = np.array(list(bdict.values()), dtype=np.float64)
  assert (len(rate_array) == len(val_array))
  return np.sum(rate_array * val_array)


def get_historical_capital(balance_df, rate_dict):
  capital_df = balance_df
  res = []
  balances = list(capital_df['balance'])
  dates = list(capital_df['trading_date'])
  for idx in range(len(balances)):
    res.append(calculate_balance_in_usd(balances[idx], rate_dict[dates[idx]]))
  capital_df['capital'] = res
  capital_df['mea'] = capital_df[['market_type', 'exchange',
                                  'api_version']].astype(str).agg('.'.join, axis=1)
  return capital_df[['strategy_name', 'strategy_group', 'mea', 'trading_date', 'capital']]


def get_gathered_data(df_groups, rolling_days):
  data = []
  for strat in df_groups.groups:
    df = df_groups.get_group(strat).dropna()
    volume_df = df[['trading_date', 'volume']]
    cum_pnl_df = df[['trading_date', 'cum_pnl']]
    df['ret_bp'] = df['pnl'].rolling(rolling_days, min_periods=1).sum().astype(
        'float') / df['volume'].rolling(rolling_days, min_periods=1).sum().astype('float') * 1e4
    retbp_df = df[['trading_date', 'ret_bp']]
    if cum_pnl_df.empty or (np.all(np.array(cum_pnl_df['cum_pnl']) == 0) and np.all(
        np.array(volume_df['volume']) == 0)):
      continue
    start_date = pd.to_datetime(FLAGS.start_date)
    volume_df = volume_df.loc[volume_df['trading_date'] >= start_date]
    cum_pnl_df = cum_pnl_df.loc[cum_pnl_df['trading_date'] >= start_date]
    retbp_df = retbp_df.loc[retbp_df['trading_date'] >= start_date]
    assert (len(volume_df) == len(cum_pnl_df) == len(retbp_df))
    if volume_df.empty and cum_pnl_df.empty and retbp_df.empty:
      continue
    data.append([strat, volume_df, cum_pnl_df, retbp_df])
  return data


def filter_strat_input(names, df):
  strat_names = list(set(df['strategy_name']))
  strat_group_names = list(set(df['strategy_group']))
  strats = []
  strat_groups = []
  for name in names:
    if name in strat_names:
      strats.append(name)
    elif name in strat_group_names:
      strat_groups.append(name)
    else:
      warnings.warn("Invalid input: %s" % name)
  return strats, strat_groups


def autolabel_bar(rects, ax, data_lst, color):
  for idx, rect in enumerate(rects):
    height = rect.get_height()
    ax.text(rect.get_x() + rect.get_width() / 2.,
            1.05 * height,
            f'{data_lst[idx]/1e6:,.3f}m',
            ha='center',
            va='bottom',
            color=color)


def autolabel_plot(ax, df, col, color, scale):
  df = df[~df.isin([np.nan, np.inf, -np.inf]).any(axis=1)]
  for row in df.iterrows():
    text = row[1][col] / scale
    ax.annotate(
        f'{text:,.3f}',
        xy=(
            row[1]['trading_date'],
            row[1][col]),
        xycoords='data',
        color=color)


def plot_pnl_volume_graph(data, figname, tmp_output_dir, days):
  fig1, ax1 = plt.subplots(len(data),
                           figsize=(max(12, min(max(12, 1.2 * days), 16)), 5 * len(data)))
  if len(data) == 1:
    ax2 = ax1.twinx()
    ax3 = ax1.twinx()
    ax1 = [ax1]
    ax2 = [ax2]
    ax3 = [ax3]
  else:
    ax2 = [ax.twinx() for ax in ax1]
    ax3 = [ax.twinx() for ax in ax1]
  fig1.suptitle(figname, y=1)
  width = min(max(24 / days, 0.5), 0.8)
  for idx, (strat, volume_df, cum_pnl_df, retbp_df) in enumerate(data):
    ax1[idx].set_ylabel("Volume")
    ax2[idx].set_ylabel("PnL")
    ax3[idx].set_ylabel("Ret_bp")
    p1 = ax1[idx].bar(
        volume_df['trading_date'],
        volume_df['volume'],
        width=width,
        color='blue',
        label='Volume',
        alpha=0.5)
    p2, = ax2[idx].plot(
        cum_pnl_df['trading_date'],
        cum_pnl_df['cum_pnl'],
        color='red',
        label='PnL',
        marker='.')
    ax3[idx].spines['right'].set_position(('outward', 60))
    p3, = ax3[idx].plot(
        retbp_df['trading_date'],
        retbp_df['ret_bp'],
        color='green',
        label='ret_bp',
        marker='*')

    lns = [p1, p2, p3]
    ax1[idx].legend(handles=lns, loc='best')

    if FLAGS.annotation:
      autolabel_bar(p1, ax1[idx], list(volume_df['volume']), 'blue')
      autolabel_plot(ax2[idx], cum_pnl_df, 'cum_pnl', 'red', 1e3)
      autolabel_plot(ax3[idx], retbp_df, 'ret_bp', 'green', 1)

    ax1[idx].yaxis.label.set_color('blue')
    ax2[idx].yaxis.label.set_color('red')
    ax3[idx].yaxis.label.set_color('green')

    ax1[idx].xaxis.set_major_locator(mdates.DayLocator())
    ax1[idx].xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
    plt.setp(ax1[idx].get_xticklabels(), rotation=45, horizontalalignment='right')
    ax3[idx].set_title(f'{strat}', y=-0.3)

  figname = os.path.join(tmp_output_dir, f'{figname}.png')
  plt.tight_layout()
  fig1.savefig(figname)
  return figname


def plot_stats(capital_df, period, tmp_output_dir):
  start_date = datetime.datetime.strptime(FLAGS.start_date, '%Y%m%d')
  end_date = datetime.datetime.strptime(FLAGS.end_date, '%Y%m%d')
  days = np.ceil((end_date - start_date).days)

  if FLAGS.plot_group_name:
    capital_df = capital_df.loc[capital_df['strategy_group'] == FLAGS.plot_group_name]
  if FLAGS.plot_selected:
    names = FLAGS.plot_selected.split(',')
    strats, strat_groups = filter_strat_input(names, capital_df)
    data = []
    if len(strats) > 0:
      capital_df['strat_list'] = strats * len(capital_df)
      capital_df['flag'] = capital_df.apply(lambda x: x['strategy_name'] in x['strat_list'], axis=1)
      capital_df_groups = capital_df.loc[capital_df['flag']].groupby('strategy_name')
      data.extend(get_gathered_data(capital_df_groups, FLAGS.rolling_days))
    if len(strat_groups) > 0:
      capital_df['strat_grp_list'] = strat_groups * len(capital_df)
      capital_df['flag'] = capital_df.apply(
          lambda x: x['strategy_group'] in x['strat_grp_list'], axis=1)
      capital_df_groups = capital_df.loc[capital_df['flag']].groupby('strategy_group').apply(
          lambda grp: grp.groupby('trading_date', as_index=False).sum())
      data.extend(
          get_gathered_data(
              capital_df_groups.groupby('strategy_group'),
              FLAGS.rolling_days))
  elif FLAGS.plot_by_group:
    capital_df_groups = capital_df.groupby('strategy_group').apply(
        lambda grp: grp.groupby('trading_date', as_index=False).sum(numeric_only=True))
    data = get_gathered_data(capital_df_groups.groupby('strategy_group'), FLAGS.rolling_days)
  else:
    capital_df_groups = capital_df.groupby('strategy_name')
    data = get_gathered_data(capital_df_groups, FLAGS.rolling_days)

  figs = []
  for i, block in enumerate(np.array_split(np.array(data, dtype=object), np.ceil(len(data) / 10))):
    try:
      figs.append(
          plot_pnl_volume_graph(
              block,
              f'pnl_volume_{period}_rolling{FLAGS.rolling_days}_bygroup{int(FLAGS.plot_by_group)}_{i}',
              tmp_output_dir,
              days))
    except Exception:
      traceback.print_exc()
  return figs


def fetch_capital_stats(tmp_output_dir):
  real_start_date = datetime.datetime.strftime(
      datetime.datetime.strptime(
          FLAGS.start_date,
          '%Y%m%d') -
      datetime.timedelta(
          days=FLAGS.rolling_days),
      '%Y%m%d')
  merged_balance_df = get_historical_account_history(real_start_date, FLAGS.end_date)
  rate_dict = get_rate_per_symbol_day(merged_balance_df, real_start_date, FLAGS.end_date)
  capital_df = get_historical_capital(merged_balance_df, rate_dict)
  pta_df = get_historical_pta(real_start_date, FLAGS.end_date)
  pta_df['trading_date'] = pta_df['trading_date'].astype('datetime64[ns]')
  capital_df['trading_date'] = capital_df['trading_date'].astype('datetime64[ns]')
  capital_df = pd.merge(capital_df, pta_df, on=['strategy_name', 'strategy_group', 'trading_date'])

  period = '-'.join([real_start_date, FLAGS.end_date])
  filename = f'capital_per_strat_{period}.csv'
  capital_df.to_csv(open(os.path.join(tmp_output_dir, filename), 'w'), index=False)
  filepath = os.path.join(tmp_output_dir, filename)
  capital_df[['pnl', 'volume']] = capital_df[['pnl', 'volume']].astype('float')
  csvtitle = f'Strat details {period}'
  return capital_df, csvtitle, filepath


def create_group_dict(df):
  pairs = [",".join(x) for x in df[['strategy_name', 'strategy_group']].values.tolist()]
  return {x.split(',')[0]: x.split(',')[1] for x in set(pairs)}


def plot_strategy_stats(capital_df, tmp_output_dir):
  period = f'{FLAGS.start_date}-{FLAGS.end_date}'
  # at most one row per day per strategy
  group_dict = create_group_dict(capital_df)
  capital_df = capital_df.groupby(['trading_date', 'strategy_name'], as_index=False).sum(numeric_only=True)
  capital_df['strategy_group'] = capital_df['strategy_name'].apply(lambda x: group_dict[x])
  if FLAGS.strategy_group:
    capital_df = capital_df.loc[capital_df['strategy_group'].isin(FLAGS.strategy_group.split(','))]
  if FLAGS.exclude_strategy:
    capital_df = capital_df.loc[~capital_df['strategy_name'].isin(
        FLAGS.exclude_strategy.split(','))]
  capital_df['trading_date'] = pd.to_datetime(capital_df['trading_date'])
  for strat in list(set(capital_df['strategy_name'])):
    cond1 = capital_df['strategy_name'] == strat
    cond2 = capital_df['trading_date'] >= pd.to_datetime(FLAGS.start_date)
    capital_df.loc[cond1 & cond2, 'cum_pnl'] = np.cumsum(capital_df.loc[cond1 & cond2, 'pnl'])
  figs = plot_stats(capital_df, period, tmp_output_dir)
  figname = f'Pnl Volume Plot {period}'
  return figname, figs


def main(argv):
  tmp_output_dir = os.path.join('/tmp', tempfile.mkdtemp())
  print(f'Results are save under {tmp_output_dir}')
  capital_df, csvtitle, filepath = fetch_capital_stats(tmp_output_dir)

  sh = None
  if FLAGS.report_slack:
    try:
      sh = SlackHelper('.', FLAGS.slack_channel)
      sh.send_slack_noti(f'{filepath} done', None, True)
    except BaseException:
      traceback.print_exc()
  try:
    send_mail_with_attachments(csvtitle, FLAGS.from_mail, FLAGS.to_mail, csvtitle, filepath, filepath)
  except BaseException:
    traceback.print_exc()

  figtitle, figs = plot_strategy_stats(capital_df, tmp_output_dir)
  if sh:
    try:
      for fig in figs:
        sh.send_slack_noti(fig, fig)
    except BaseException:
      traceback.print_exc()
  try:
    send_mail_with_attachments(figtitle, FLAGS.from_mail, FLAGS.to_mail, figtitle, '', figs)
  except BaseException:
    traceback.print_exc()


"""
./pyrunner python/coin/strategy/basis_strat/analysis/report_live_strat_capital.py --start_date=20220125 --end_date=20220325 --rolling_days=5 --plot_group_name=basis_smm2
./pyrunner python/coin/strategy/basis_strat/analysis/report_live_strat_capital.py --start_date=20220125 --end_date=20220325 --rolling_days=5 --plot_by_group=True
./pyrunner python/coin/strategy/basis_strat/analysis/report_live_strat_capital.py --start_date=20220125 --end_date=20220325 --rolling_days=5 --plot_selected='basis_smm2_binance_spot_binance_perpetual,basis_smm2'
"""


if __name__ == '__main__':
  flags.DEFINE_string('start_date', '20220101', '')
  flags.DEFINE_string('end_date', '20220301', '')
  flags.DEFINE_string('from_mail', 'basis_strategy@joomo.io', '')
  flags.DEFINE_string('to_mail', 'zijin@prestolabs.io', '')
  flags.DEFINE_bool('report_slack', True, '')
  flags.DEFINE_string('slack_channel', 'team_u_monitor', '')
  flags.DEFINE_bool('cache_clear', False, '')
  flags.DEFINE_integer('max_workers', 14, '')
  # plot params
  flags.DEFINE_integer('chunk_size', 1, '')
  flags.DEFINE_integer('rolling_days', 5, '')
  flags.DEFINE_bool('annotation', True, 'Show the value annotation on plot')
  flags.DEFINE_string('strategy_group', '', 'a list of group to include')
  flags.DEFINE_string('exclude_strategy', '', 'a list of strategy name to exclude')
  flags.DEFINE_string('plot_group_name', '', 'Specify strategies in which strategy group to plot')
  flags.DEFINE_bool(
      'plot_by_group',
      False,
      'If false, plot by each strat. If true, plot by each strat group.')
  flags.DEFINE_string(
      'plot_selected',
      '',
      'plot only specified strategies/strategy_groups, split by comma')
  app.run(main)
